#! -*- coding: utf-8 -*-
"""
@Info: 从标注的excel文件中转换为训练txt数据
"""
import pandas as pd
import re
import json

df = pd.read_excel('../../data/AIPS语料-20240716.xlsx')
# df = df[df.iloc[:, 7] == 2]  # 筛选已标注用于训练的数据
df = df[df.iloc[:, 7] != 10]  # 这部分语料不清晰，需确认
df.rename(columns={df.columns[4]: 'operate_intent'}, inplace=True)
df.rename(columns={df.columns[5]: 'text'}, inplace=True)
df.rename(columns={df.columns[6]: 'label'}, inplace=True)

df = df.dropna(subset=['text', 'operate_intent', 'label'], how='any')
df = df.iloc[1:, [5, 6, 4]]
df['operate_intent'] = df['operate_intent'].map(str)

schema = []
# print('operate_intent:', set(df['operate_intent']))
# intent_schema = {'查询': set(), '插单': set(), '重排': set()}
intent_schema = {}

OPERATE_TYPE = [str(i) for i in [1, 2, 3, 4, 5, 6, 9, 11,
                                 12, 13, 14, 15, 16, 17, 19, 20, 21, 22]]
for i in OPERATE_TYPE:
    intent_schema.update({i: set()})

df['cats'] = ''
df['entities'] = None
for index, row in df.iterrows():
    matches = re.findall(r"\[E[^\]]+\]", row['label'])
    for match in matches:
        matched_schema = match.replace('E', '').replace('[', '').replace(']', '')
        schema.append(matched_schema)
        intent_schema[row['operate_intent']].add(matched_schema)
        df.loc[index, 'cats'] = [matched_schema]

    matches = re.findall(r"\{[^\}E]+\}", row['label'])
    entities = []
    for match in matches:
        ner_string = re.sub(r'\[.*?\]', '', match.replace('{', '').replace('}', ''))
        start = df.loc[index, 'text'].find(ner_string)
        # 结束位置就是字符结束的位置
        end = start + len(ner_string)
        entities.append([start, end, re.search(r'\[.*?\]', match).group(0).replace('[', '').replace(']', '')])
    df.loc[index, 'entities'] = json.dumps(entities, ensure_ascii=False)

print(set(schema))  # 所有schema
intent_schema_list = []
for key, value in intent_schema.items():
    intent_schema_list.append({'operate_type': key, 'schema': list(value)})
print(intent_schema_list)

# train_df, eval_df = train_test_split(df, test_size=0.1, random_state=42)

df.drop(['label', 'operate_intent'], axis=1, inplace=True) #删除label和operate_intent列
train_df = df
eval_df = df

# 将结果写入文件
with open('../data/all.jsonl', 'w', encoding='utf-8') as json_file:
    for _, row in train_df.iterrows():
        data = row.to_dict()
        data['entities'] = json.loads(data['entities'])
        json_file.write(json.dumps(data, ensure_ascii=False) + '\n')

with open('../data/all_test.jsonl', 'w', encoding='utf-8') as json_file:
    for _, row in eval_df.iterrows():
        data = row.to_dict()
        data['entities'] = json.loads(data['entities'])
        json_file.write(json.dumps(data, ensure_ascii=False) + '\n')

with open('../src/schema.txt', 'w', encoding='utf-8') as f:
    f.write(json.dumps(intent_schema_list, ensure_ascii=False))
